import random

import matplotlib.pyplot as plt
import numpy as np
from scipy.stats import multivariate_normal
from sklearn.datasets import load_iris


def generate_data(sample_num: [], u: [[]], sigma: [[[]]]) -> (np.ndarray, np.ndarray):
    """
    根据u,sigma产生样本len(sample_num)类数据，其中每类数据的个数为sample_num中的每个元素
    u[i]表示该第i类的(x,y)均值
    sigma[i]表示该第i类的(x,y)协方差矩阵

    :param sample_num: 列表，len(sample_num)=num_classes，每类样本个数
    :param u: shape=(num_classes,2)的样本均值均值矩阵
    :param sigma: shape=(num_classes,2,2)的协方差矩阵
    :return:x_train:shape=(sum(sample_num),2), y_train(sum(sample_num),1)
    """
    total = sum(sample_num)
    x_train = np.zeros((total, 2), dtype=np.float32)
    y_train = np.zeros((total, 1), dtype=np.int32)
    x_train[:sample_num[0], :] = np.random.multivariate_normal(u[0], sigma[0], size=sample_num[0])
    y_train[:sample_num[0]] = 0
    base = 0
    for i in range(len(sample_num)):
        x_train[base:base + sample_num[i], :] = np.random.multivariate_normal(u[i], sigma[i], size=sample_num[i])
        y_train[base:base + sample_num[i]] = i
        base += sample_num[i]
    return x_train, y_train


def draw_picture_2d(x_train: np.ndarray, y_train: np.ndarray, title: str = 'picture'):
    """
    绘制二维散点图以及决策函数

    :param title: 标题名称
    :param x_train: 二维向量
    :param y_train: 一维向量，与x_train元素数目相同
    :return: 无
    """
    x_min = min(x_train[:, 0]) - 1
    x_max = max(x_train[:, 0]) + 1
    y_min = min(x_train[:, 1]) - 1
    y_max = max(x_train[:, 1]) + 1
    plt.axis([x_min, x_max, y_min, y_max])  # 规定坐标轴的值
    plt.scatter(x=x_train[:, 0], y=x_train[:, 1], c=list(y_train), marker='o')
    plt.title(title)
    plt.ylabel('y')
    plt.xlabel('x')
    plt.show()


def k_means(x_train: np.ndarray, num_classes: int, random_state: int = 1) -> np.ndarray:
    """
    k_means算法

    :param x_train: shape=(sample_num,dimension)的待分类样本
    :param num_classes: 分类数目
    :param random_state: 随机种子
    :return: y_pred: shape=(sample_num,1)的已分类标签
    """
    # 获取初始样本中心（随机）
    total = x_train.shape[0]
    dimension = x_train.shape[1]
    y_train = np.zeros((total, 1), dtype=np.int32)
    base = 0
    for i in range(num_classes):
        y_train[base:base + int(total / num_classes)] = i
        base += int(total / num_classes)
    y_train[base:] = 0
    index = [i for i in range(total)]
    random.seed(random_state)
    random.shuffle(index)
    y_train = y_train[index]
    # 开始迭代
    variation = 100
    while variation != 0:
        variation = 0
        # M步：根据标签，计算k个聚类中心（centers）
        centers = np.zeros((num_classes, dimension), dtype=np.float32)
        centers_num = np.zeros((num_classes, 1), dtype=np.int32)
        for i in range(total):
            for j in range(num_classes):
                if y_train[i] == j:
                    centers_num[j] += 1
                    centers[j] += x_train[i]
                    break
        for j in range(num_classes):
            centers[j] = centers[j] / centers_num[j]
        # E步：遍历所有样本，根据样本到k个聚类中心的距离度量，更新该样本的标签
        for i in range(total):
            min_j = 0
            min_distance = np.linalg.norm(x_train[i] - centers[0])
            for j in range(1, num_classes):
                distance = np.linalg.norm(x_train[i] - centers[j])
                if distance < min_distance:
                    min_j = j
                    min_distance = distance
            if y_train[i] != min_j:
                variation += 1
                y_train[i] = min_j
    return y_train


def gaussian_mixed_model(x_train: np.ndarray, num_classes: int, random_state: int = 1, epsilon: float = 1e-5,
                         epochs: int = 500) -> np.ndarray:
    """
    gaussian_mixed_model算法

    :param x_train: shape=(sample_num,dimension)的待分类样本
    :param num_classes: 分类数目
    :param random_state: 随机种子
    :param epsilon: 似然值的变化小于epsilon时迭代停止
    :param epochs: 迭代次数大于epochs时迭代停止
    :return: y_pred: shape=(sample_num,1)的已分类标签
    """
    np.random.seed(random_state)
    N = x_train.shape[0]  # 样本数
    dimension = x_train.shape[1]  # 每个样本的维度数
    K = num_classes  # 类别数
    gama = np.random.rand(N, K)
    likelihood_pre = -1
    likelihood_now = 0
    epochs_now = 0
    while abs(likelihood_now - likelihood_pre) > epsilon and epochs_now < epochs:  # 似然值的变化大于epsilon时迭代继续
        likelihood_pre = likelihood_now
        epochs_now += 1
        # M步：更新u,sigma,pi
        u = np.zeros(shape=(K, dimension))
        sigma = np.zeros(shape=(K, dimension, dimension))
        pi = np.zeros(shape=(K, 1))
        for k in range(K):
            n_k = 0
            temp_u = 0
            for n in range(N):
                n_k += gama[n, k]
                temp_u += gama[n, k] * x_train[n]
            u[k] = temp_u / n_k
            temp_sigma = np.zeros(shape=(dimension, dimension), dtype=np.float64)
            for n in range(N):
                temp_distance = np.array([x_train[n] - u[k]])
                temp_sigma += gama[n, k] * temp_distance.T.dot(temp_distance)
            sigma[k] = temp_sigma / n_k
            pi[k] = n_k / N
        # E步：更新gama
        for n in range(N):
            temp_gama = 0
            for k in range(K):
                temp_gama += pi[k] * multivariate_normal.pdf(x_train[n], mean=u[k], cov=sigma[k])
            for k in range(K):
                gama[n, k] = pi[k] * multivariate_normal.pdf(x_train[n], mean=u[k], cov=sigma[k]) / temp_gama
        # 重新计算似然值
        likelihood_now = 0
        for n in range(N):
            temp_likelihood = 0
            for k in range(K):
                temp_likelihood += pi[k] * multivariate_normal.pdf(x_train[n], mean=u[k], cov=sigma[k])
            likelihood_now += np.emath.log(temp_likelihood)
        if epochs_now % 10 == 0:
            print('epoch：{}\t似然值：{}'.format(epochs_now, likelihood_now))
            yield epochs_now, likelihood_now, np.argmax(gama, axis=-1)
    yield epochs_now, likelihood_now, np.argmax(gama, axis=-1)


def k_means_instance():
    global random_state, _, y_pred
    print('\033[0;34;40-m---k_means---\033[0m')
    random_state = 1
    for _ in range(4):
        y_pred = k_means(x_train, 3, random_state=random_state)
        random_state += 1
        draw_picture_2d(x_train, y_pred, 'cluster[random_state={}]'.format(random_state))


def gaussian_mixed_model_instance():
    global random_state, _, y_pred
    print('\033[0;34;40-m---gaussian_mixed_model---\033[0m')
    random_state = 1
    for _ in range(4):
        print('--------------random_state={}-----------------------'.format(random_state))
        for (epochs_now, likelihood_now, y_pred) in gaussian_mixed_model(x_train, 3, random_state=random_state):
            draw_picture_2d(x_train, y_pred,
                            'cluster[random_state={} epoch={} likelihood={}]'.format(random_state, epochs_now,
                                                                                     likelihood_now))
        random_state += 1


def iris_instance():
    global x_train, y_train, y_pred
    # 鸢尾花数据集
    iris = load_iris()
    iris_data = np.array(iris.data)
    iris_target = np.array(iris.target).T
    x_train = iris_data[:, 1:3]
    y_train = iris_target
    draw_picture_2d(x_train, y_train, 'iris_original')
    random_state = 1
    for _ in range(4):
        print('--------------random_state={}-----------------------'.format(random_state))
        for (epochs_now, likelihood_now, y_pred) in gaussian_mixed_model(x_train, 3, random_state=random_state):
            draw_picture_2d(x_train, y_pred,
                            'iris_gaussian_mixed_model[random_state={} epoch={} likelihood={}]'.format(random_state,
                                                                                                       epochs_now,
                                                                                                       likelihood_now))
        random_state += 1


if __name__ == '__main__':
    # 自己产生3个类别的数据

    sample_num = [200, 100, 50]  # 每一类样本个数
    u = [[2, 2],
         [7, 7],
         [10, 4]]  # 每类样本均值
    sigma = [[[1, 0.5], [0.5, 1]],
             [[2, 0], [0, 2]],
             [[2, 1], [1, 2]]]  # 每类样本方差

    x_train, y_train = generate_data(sample_num, u, sigma)
    # draw_picture_2d(x_train, y_train, 'original')

    # k_means
    # k_means_instance()

    # GMM
    # gaussian_mixed_model_instance()

    # GMM对鸢尾花数据集进行测试
    iris_instance()
